#!/bin/env python
# -*- coding: utf-8 -*-
import taichi as ti
import numpy as np
import cv2
from copy import deepcopy
from typing import List, Tuple, Optional, Union
from sklearn.cluster import KMeans
from scipy.spatial import cKDTree
from tqdm import tqdm
from PIL import Image
import os

@ti.data_oriented
class DitherUtil:
    """
    抖动工具类，提供图像处理和抖动算法的共用功能。
    """
    
    def __init__(self, arch=ti.cpu, debug=False, max_value=255):
        """
        初始化抖动工具类。
        
        :param arch: Taichi架构，默认为CPU，可以设置为ti.gpu启用GPU加速
        :param debug: 是否开启调试模式，默认为False
        :param max_value: 像素最大值，默认为255
        """
        # 初始化Taichi
        # 添加明确的GPU内存配置，避免默认值太小导致GPU资源分配不足
        if arch == ti.gpu:
            ti.init(arch=arch, debug=debug, device_memory_GB=2.0, default_fp=ti.f32, default_ip=ti.i32, 
                    kernel_profiler=True, print_ir=debug, print_kernel_llvm_ir=debug)
        else:
            ti.init(arch=arch, debug=debug)
            
        self.arch = arch
        self.DEBUG = debug
        self.MAX = max_value
        np.set_printoptions(threshold=np.inf, linewidth=180)  # numpy打印选项
        
        # 定义类型
        self.type_img2d = ti.types.ndarray(element_dim=0, ndim=2)
        self.type_bayerM = ti.types.ndarray(element_dim=0, ndim=2)
        self.type_img_color = ti.types.ndarray(element_dim=1, ndim=2)  # 彩色图像类型
        
        # 默认调色板（黑白）
        self.PALETTE = [0x00, 0xFF]
        self.color_mode = 'grayscale'  # 'grayscale' 或 'color'
        self.color_palette = []  # 彩色调色板 [(r,g,b), ...]
        
        # 如果使用GPU，打印GPU设备信息
        if arch == ti.gpu:
            self.print_gpu_info()
    
    def print_gpu_info(self):
        """
        打印GPU设备信息
        """
        print("===== GPU 信息 =====")
        try:
            # 获取当前Taichi后端信息
            import taichi as ti
            backend = ti.lang.impl.current_cfg().arch
            
            if backend == ti.cuda:
                print(f"使用 CUDA 后端")
            elif backend == ti.vulkan:
                print(f"使用 Vulkan 后端")
            elif backend == ti.opengl:
                print(f"使用 OpenGL 后端")
            else:
                print(f"使用其他GPU后端")
        except Exception as e:
            print(f"无法获取GPU详细信息: {e}")
            print(f"Taichi版本: {ti.__version__}")
            
        print("===================")

    def resize_image(self, image, width=None, height=None):
        """
        调整图像大小，保持宽高比。

        :param image: 输入图像
        :param width: 目标宽度，如果为None则根据height等比例计算
        :param height: 目标高度，如果为None则根据width等比例计算
        :return: 调整大小后的图像
        """
        # 获取原始尺寸
        h, w = image.shape[:2]
        
        # 如果宽度和高度都未指定，返回原图
        if width is None and height is None:
            return image
        
        # 计算缩放比例
        if width is None:
            # 根据高度等比例缩放
            ratio = height / float(h)
            new_width = int(w * ratio)
            new_height = height
        elif height is None:
            # 根据宽度等比例缩放
            ratio = width / float(w)
            new_width = width
            new_height = int(h * ratio)
        else:
            # 宽度和高度都指定了
            new_width = width
            new_height = height
        
        # 调整大小
        resized = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA)
        
        if self.DEBUG:
            print(f"图像已调整大小: {w}x{h} -> {new_width}x{new_height}")
        
        return resized
    
    def set_palette(self, palette: List[int]):
        """
        设置用于抖动的调色板。
        
        :param palette: 包含色值的列表
        """
        self.PALETTE = palette
    
    def generate_palette(self, num_colors: int):
        """
        生成具有指定数量色值的调色板。
        
        :param num_colors: 需要的色值数量
        :return: 生成的调色板列表
        """
        if num_colors < 2:
            num_colors = 2  # 至少需要两种颜色
        
        # 计算均匀分布的色值
        palette = [int(i * self.MAX / (num_colors - 1)) for i in range(num_colors)]
        self.PALETTE = palette
        return palette
    
    def generate_color_palette_from_image(self, image, num_colors: int):
        """
        从输入图像中提取最具代表性的颜色作为调色板。
        使用K-Means聚类算法来找到图像中最主要的颜色。
        
        :param image: 输入彩色图像
        :param num_colors: 需要提取的颜色数量
        :return: 颜色调色板，格式为 [(r,g,b), ...]
        """
        # 确保输入是彩色图像
        if len(image.shape) < 3:
            raise ValueError("需要输入彩色图像来提取调色板")
        
        # 将图像重塑为像素列表
        pixels = image.reshape(-1, 3).astype(np.float32)
        
        # 随机抽样以加速处理（对于大图像）
        sample_size = min(100000, pixels.shape[0])
        indices = np.random.choice(pixels.shape[0], sample_size, replace=False)
        pixel_samples = pixels[indices]
        
        # 使用K-Means聚类
        kmeans = KMeans(n_clusters=num_colors, random_state=0, n_init=1)
        kmeans.fit(pixel_samples)
        
        # 获取聚类中心作为调色板颜色
        colors = kmeans.cluster_centers_.astype(np.uint8)
        
        # 创建RGB格式的调色板
        palette = [tuple(color) for color in colors]
        self.color_palette = palette
        return palette
    
    def show_image(self, img):
        """
        显示图像。
        
        :param img: 图像数据或图像文件路径
        """
        if isinstance(img, str):
            img = cv2.imread(img)
        elif img.dtype != 'uint8':
            img = np.clip(img, 0, self.MAX)  # 将图片像素值限制在 0~MAX 之间
            img = img.astype(np.uint8)
        
        # 显示图片
        cv2.imshow('Image', img)
        while True:
            key = cv2.waitKey(1) & 0xFF
            if key == 27:  # 27 是 ESC 键的 ASCII 码
                break
        cv2.destroyAllWindows()
    
    @ti.func
    def clamp(self, x, min=-255, max=255):
        """
        控制出血阈值，限制值的范围。
        
        :param x: 要限制的值
        :param min: 最小值，建议在 [-256,0] 范围内
        :param max: 最大值
        :return: 限制后的值
        """
        return ti.math.clamp(x, min, max)
        
    # Taichi优化的彩色图像处理函数
    @ti.kernel
    def prepare_color_img(self, color_img: ti.types.ndarray(), gray_img: ti.types.ndarray()):
        """
        将彩色图像转换为灰度图像，以便Taichi处理
        
        :param color_img: 输入彩色图像
        :param gray_img: 输出灰度图像
        """
        h, w = gray_img.shape
        for i, j in ti.ndrange(h, w):
            # 使用加权平均法计算灰度值
            r = color_img[i, j, 0]
            g = color_img[i, j, 1]
            b = color_img[i, j, 2]
            gray_img[i, j] = ti.cast(0.299 * r + 0.587 * g + 0.114 * b, ti.i32)

    # 优化的颜色匹配函数
    def find_best_color_match(self, pixel, palette):
        """
        找到调色板中最接近给定像素的颜色。
        使用numpy向量化操作加速。
        
        :param pixel: RGB像素值，形状为(3,)
        :param palette: 颜色调色板，形状为(n, 3)
        :return: 最佳匹配的颜色索引和颜色值
        """
        # 计算所有颜色的欧氏距离
        distances = np.sum((palette - pixel) ** 2, axis=1)
        # 找到最小距离的索引
        best_idx = np.argmin(distances)
        return best_idx, palette[best_idx]
    
    def diff(self, img1, img2):
        """
        计算两个图像之间的差异。
        
        :param img1: 第一个图像
        :param img2: 第二个图像
        :return: 差异值
        """
        return np.sum(np.abs(img1-img2))
    
    def save_image(self, img, filename):
        """
        保存图像到文件。
        
        :param img: 要保存的图像
        :param filename: 保存的文件名
        """
        if img.dtype != 'uint8':
            img = np.clip(img, 0, self.MAX)
            img = img.astype(np.uint8)
        
        cv2.imwrite(filename, img)
        if self.DEBUG:
            print(f"图像已保存到: {filename}")
            
    def save_bmp_with_palette(self, img, filename, bit_depth=8):
        """
        保存带有调色板和指定位深度的BMP图像。
        使用抖动后图像中实际使用的颜色作为调色板。
        
        :param img: 要保存的图像
        :param filename: 保存的文件名
        :param bit_depth: 位深度，支持1、4、8位，默认为8
        """
        # 确保图像是uint8类型
        if img.dtype != 'uint8':
            img = np.clip(img, 0, self.MAX)
            img = img.astype(np.uint8)
            
        # 根据图像的维度确定保存模式
        if len(img.shape) == 2:
            # 灰度图像
            pil_img = Image.fromarray(img)
            pil_img.save(filename, format="BMP", bits=bit_depth)
            
            if self.DEBUG:
                print(f"BMP图像已保存到: {filename}")
            return
        
        elif len(img.shape) == 3 and img.shape[2] == 3:
            # 彩色图像
            # OpenCV使用BGR顺序，而PIL使用RGB顺序，所以需要转换
            img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            
            # 创建PIL图像
            pil_img = Image.fromarray(img_rgb)
            
            # 提取实际使用的颜色
            # 使用PIL的quantize方法获取更精确的调色板
            max_colors = 2**bit_depth
            
            # 直接使用原始图像的颜色量化
            pil_img_quantized = pil_img.quantize(colors=max_colors, method=Image.MEDIANCUT)
            
            # 从量化后的图像中获取调色板
            palette = pil_img_quantized.getpalette()
            
            # 获取调色板对应的索引图像
            index_img = np.array(pil_img_quantized)
            
            # 创建空的PIL图像为带索引的图像
            output_img = Image.new('P', (img.shape[1], img.shape[0]))
            
            # 设置调色板
            output_img.putpalette(palette)
            
            # 用索引填充图像
            output_img.paste(pil_img_quantized)
            
            # 保存为BMP
            output_img.save(filename, format="BMP", bits=bit_depth)
            
            if self.DEBUG:
                num_colors = len(set(index_img.flatten()))
                print(f"带有{num_colors}色调色板的{bit_depth}位BMP图像已保存到: {filename}")
            return
        
        # 如果无法使用调色板（例如位深度不支持或没有调色板），则使用常规方式保存
        if self.DEBUG:
            print(f"无法以{bit_depth}位深度保存带调色板的图像，将使用默认保存方式")
        self.save_image(img, filename)

    # 获取系统信息
    def get_system_info(self):
        """
        获取系统信息，包括Taichi配置和设备信息
        
        :return: 系统信息字符串
        """
        info = []
        info.append(f"Taichi version: {ti.__version__}")
        info.append(f"Arch: {self.arch}")
        info.append(f"Debug mode: {self.DEBUG}")
        
        # 获取平台信息
        import platform
        info.append(f"Platform: {platform.platform()}")
        info.append(f"Python: {platform.python_version()}")
        
        # 如果使用GPU，尝试获取GPU信息
        if self.arch == ti.gpu:
            try:
                # 尝试使用GPU信息库
                import GPUtil
                gpus = GPUtil.getGPUs()
                if gpus:
                    for i, gpu in enumerate(gpus):
                        info.append(f"GPU {i}: {gpu.name} (Memory: {gpu.memoryTotal}MB)")
            except ImportError:
                info.append("GPU information unavailable (GPUtil not installed)")
        
        return "\n".join(info) 

    @ti.kernel
    def map_to_palette(self, img: ti.types.ndarray(), palette: ti.types.ndarray()):
        """
        使用Taichi内核将图像像素映射到最接近的调色板颜色。
        
        :param img: 输入/输出图像数组
        :param palette: 调色板数组
        """
        h, w = img.shape
        n_colors = palette.shape[0]
        
        for i, j in ti.ndrange(h, w):
            pixel_value = img[i, j]
            min_dist = abs(pixel_value - palette[0])
            best_idx = 0
            
            # 寻找最接近的调色板颜色
            for k in range(1, n_colors):
                dist = abs(pixel_value - palette[k])
                if dist < min_dist:
                    min_dist = dist
                    best_idx = k
            
            # 应用最接近的调色板颜色
            img[i, j] = palette[best_idx] 

    @ti.kernel
    def map_to_color_palette(self, img: ti.types.ndarray(), palette: ti.types.ndarray()):
        """
        使用Taichi内核将彩色图像像素映射到最接近的调色板颜色。
        
        :param img: 输入/输出彩色图像数组 (h, w, 3)
        :param palette: 调色板数组 (n, 3)
        """
        h, w = img.shape[0], img.shape[1]
        n_colors = palette.shape[0]
        
        for i, j in ti.ndrange(h, w):
            r = img[i, j, 0]
            g = img[i, j, 1]
            b = img[i, j, 2]
            
            min_dist = 1e10
            best_idx = 0
            
            # 寻找最接近的调色板颜色
            for k in range(n_colors):
                pr = palette[k, 0]
                pg = palette[k, 1]
                pb = palette[k, 2]
                
                # 计算颜色距离
                dist = (r-pr)*(r-pr) + (g-pg)*(g-pg) + (b-pb)*(b-pb)
                
                if dist < min_dist:
                    min_dist = dist
                    best_idx = k
            
            # 应用最接近的调色板颜色
            img[i, j, 0] = palette[best_idx, 0]
            img[i, j, 1] = palette[best_idx, 1]
            img[i, j, 2] = palette[best_idx, 2] 